import os
import argparse
from hashlib import pbkdf2_hmac

from Cryptodome.Cipher import AES


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "-f",
        "--file",
        type=str,
        required=True,
        help="Input file.",
    )
    parser.add_argument(
        "-o",
        "--out",
        type=str,
        required=False,
        help="Output file.",
    )
    parser.add_argument(
        "-p",
        "--password",
        type=str,
        required=True,
        help="Password.",
    )
    return parser.parse_args()


def main():
    args = parse_args()

    if not os.path.exists(args.file):
        print("* Invalid filepath")
        return

    with open(args.file, "rb") as file:
        salt, nonce, tag, ciphertext = [
            file.read(offset) for offset in (16, 16, 16, -1)
        ]

    dervied_key = pbkdf2_hmac(
        "sha256",
        args.password.encode("utf-8"),
        salt,
        500_000,
    )

    cipher = AES.new(dervied_key, AES.MODE_EAX, nonce)
    data = cipher.decrypt_and_verify(ciphertext, tag)

    if not args.out:
        print(data.decode("utf-8"))
        return

    with open(args.out, "wb") as file:
        file.write(data)


if __name__ == "__main__":
    main()
